import torch
from torchvision import transforms
from PIL import Image, ImageOps

import numpy as np
import scipy.misc as misc
import os
import glob
import cv2
from networks.whole_network import Mymodel
from utils.misc import thresh_OTSU, ReScaleSize, Crop
# from utils.model_eval import eval
from utils.evaluation_metrics import get_acc

DATABASE = './DRIVE/'
#
args = {
    # 'root'     : './dataset/' + DATABASE,
    # 'test_path': './dataset/' + DATABASE + 'test/',
    'test_path': '/home/jiayu/Desktop/Jimmy/1219(分六组)/601-1219_all_nerve/img/',
    'pred_path': 'assets/' + 'JiaoMo/',
    'img_size': 384
}

if not os.path.exists(args['pred_path']):
    os.makedirs(args['pred_path'])


def rescale(img):
    w, h = img.size
    min_len = min(w, h)
    new_w, new_h = min_len, min_len
    scale_w = (w - new_w) // 2
    scale_h = (h - new_h) // 2
    box = (scale_w, scale_h, scale_w + new_w, scale_h + new_h)
    img = img.crop(box)
    return img


def ReScaleSize_DRIVE(image, re_size=512):
    w, h = image.size
    min_len = min(w, h)
    new_w, new_h = min_len, min_len
    scale_w = (w - new_w) // 2
    scale_h = (h - new_h) // 2
    box = (scale_w, scale_h, scale_w + new_w, scale_h + new_h)
    image = image.crop(box)
    image = image.resize((re_size, re_size))
    return image  # , origin_w, origin_h


def ReScaleSize_STARE(image, re_size=512):
    w, h = image.size
    max_len = max(w, h)
    new_w, new_h = max_len, max_len
    delta_w = new_w - w
    delta_h = new_h - h
    padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
    image = ImageOps.expand(image, padding, fill=0)
    # origin_w, origin_h = w, h
    image = image.resize((re_size, re_size))
    return image  # , origin_w, origin_h


def load_JiaoMo():
    test_images = []
    # test_labels = []
    for file in glob.glob(os.path.join(args['test_path'],  '*')):

        # basename = os.path.basename(file)
        # img_path = os.path.join(args['test_path'], 'img', basename)
        img_path = file
        # label_path = os.path.join(args['test_path'], 'gt3', basename)
        test_images.append(img_path)
        # test_labels.append(label_path)
    return test_images


def load_net():
    # net = torch.load('./110_Our_FFN(0.8738,0.1813).pkl', map_location=torch.device('cuda:0'))
    net = Mymodel()
    weight = torch.load('./110_Our_FFN(0.8738,0.1813).pth')
    net.load_state_dict(weight)
    net = net.cuda()
    return net

#
# def save_prediction(pred, filename=''):
#     save_path = args['pred_path'] + 'pred/'
#     if not os.path.exists(save_path):
#         os.makedirs(save_path)
#         print("Make dirs success!")
#
#     mask = pred.data.cpu().numpy() * 255  # .data转换成tensor
#     mask = np.transpose(np.squeeze(mask, axis=0), [1, 2, 0])
#     mask = np.squeeze(mask, axis=-1)  # 只留下高和宽
#     cv2.imwrite(save_path + filename + '.png', mask)


def predict(image):
    net = load_net()
    # images, labels = load_nerve()
    # images = load_JiaoMo()
    # images, labels = load_stare()
    # images, labels = load_padova1()
    # images, labels = load_octa()

    transform = transforms.Compose([
        transforms.ToTensor()
    ])

    with torch.no_grad():
        net.eval()
        # for i in range(len(images)):
        # print(images[i])
        # name_list = images[i].split('/')
        # index = name_list[-1][:-4]  # 取去掉.tif（最后这四位）的basename
        # image = Image.open(images[i]).convert("L")

        image = transform(image)
        image = image.unsqueeze(0)  # 将三维提升到四维，因为网络输入要求是四维
        image = image.cuda()
        output = net(image)  # output.size(): [1, 1, 384, 384]
        output = output.data.cpu().numpy()
        output = output.squeeze(0)
        output = output.squeeze(0)  ###预测、标签和原图方向不太一致，旋转了，要改squeeze这一块
        output = (output * 255).astype(np.uint8)
        # cv2.imwrite("/home/jiayu/Desktop/Jimmy/1219(分六组)/601-1219_all_nerve/predict_old_model/" + index + ".png", output)
        # cv2.imwrite("/home/imed/文档/tmp/label/" + index + ".png", label)
        # save_prediction(output, filename=index + '_pred')
        # print("output saving successfully")

    return output

if __name__ == '__main__':
    # image = Image.open("batch_1_1.tif").convert("L")
    # result = predict(image)
    #
    # cv2.imwrite("./AAAAAAAA.png", result)

    net = load_net()
    # torch.save(net.module.state_dict(), './110_Our_FFN(0.8738,0.1813).pth')

